import sys

import pandas as pd
import numpy as np
from scipy import optimize
from scipy.optimize import linprog

# x为投入，y_g为产出，cur为当前计算的dmu，rts=1为规模报酬可变，rts=0为规模报酬不变
def sbmeff(x, y_g, cur, rts=0):
    m, n = x.shape  # m为投入变量数，n为决定单元数
    s1 = y_g.shape[0]
    f = np.concatenate([np.zeros(n), -1/(m*x[:, cur]),
                        np.zeros(s1), np.array([1])])
    Aeq1 = np.hstack([x,
                      np.identity(m),
                      np.zeros((m, s1)),
                      -x[:, cur, None]])
    Aeq2 = np.hstack([y_g,
                      np.zeros((s1, m)),
                      -np.identity(s1),
                      -y_g[:, cur, None]])
    Aeq4 = np.hstack([np.zeros(n),
                      np.zeros(m),
                      1/((s1)*(y_g[:, cur])),
                      np.array([1])]).reshape(1, -1)
    if (rts == 1):
        Aeq5 = np.hstack([np.ones(n),
                          np.zeros((m+s1)),
                          np.array([-1])]).reshape(1, -1)
        Aeq = np.vstack([Aeq1, Aeq2, Aeq4, Aeq5])
        beq = np.concatenate(
            [np.zeros(m+s1), np.array([1]), np.array([0])])
    else:
        Aeq = np.vstack([Aeq1, Aeq2, Aeq4])
        beq = np.concatenate([np.zeros(m+s1), np.array([1])])
    bounds = tuple([(0, None) for t in range(n+s1+m+1)])
    res = optimize.linprog(c=f, A_eq=Aeq, b_eq=beq, bounds=bounds)
    return res



# 不含非期望产出的超效率SBM
# x为投入，y_g为产出，cur为当前计算的dmu，rts=1为规模报酬可变，rts=0为规模报酬不变
def sup_sbmeff(x, y_g, cur, rts=0):
    m, n = x.shape  # m为投入变量数，n为决定单元数
    s1 = y_g.shape[0]
    f = np.concatenate([np.zeros(n), 1/(m*x[:, cur]),
                        np.zeros(s1), np.array([1])])
    Aeq1 = np.hstack([np.zeros(n),
                      np.zeros(m),
                      -1/((s1)*(y_g[:, cur])),
                      np.array([1])]).reshape(1, -1)
    if (rts == 1):  # 规模报酬可变情况
        Aeq2 = np.hstack([np.ones(n),
                          np.zeros((m+s1)),
                          np.array([-1])]).reshape(1, -1)  # 规模报酬可变增加的约束
        Aeq = np.vstack([Aeq1, Aeq2])  # 合并等式约束矩阵
        Aeq[:, cur] = 0  # 将正在计算列的系数置为0，对应约束条件中的求和不包括当前dmu
        beq = np.concatenate([np.array([1]), np.array([0])])  # 增加VRS约束方程右边的0
    else:
        Aeq = Aeq1
        beq = np.array([1])
    Aub1 = np.hstack([x,
                      -np.identity(m),
                      np.zeros((m, s1)),
                      -x[:, cur, None]])
    Aub2 = np.hstack([-y_g,
                      np.zeros((s1, m)),
                      -np.identity(s1),
                      y_g[:, cur, None]])
    Aub = np.vstack([Aub1, Aub2])
    bub = np.zeros(m+s1)
    Aub[:, cur] = 0 # 将正在计算列的系数置为0，对应约束条件中的求和不包括当前dmu
    bounds = tuple([(0, None) for i in range(n+s1+m+1)])
    res = optimize.linprog(c=f, A_ub=Aub, b_ub=bub,
                           A_eq=Aeq, b_eq=beq, bounds=bounds)

    return res


# 不含非期望产出的超效率SBM
# x为投入，y_g为产出，y_b为非期望产出，cur为当前计算的dmu，rts=1为规模报酬可变，rts=0为规模报酬不变
def un_sbmeff(x, y_g, y_b, cur, rts=0):
    m, n = x.shape  # m为投入变量数，n为决定单元数
    s1 = y_g.shape[0]
    s2 = y_b.shape[0]
    f = np.concatenate([np.zeros(n), -1/(m*x[:, cur]),
                        np.zeros(s1+s2), np.array([1])])
    Aeq1 = np.hstack([x,
                      np.identity(m),
                      np.zeros((m, s1+s2)),
                      -x[:, cur, None]])
    Aeq2 = np.hstack([y_g,
                      np.zeros((s1, m)),
                      -np.identity(s1),
                      np.zeros((s1, s2)),
                      -y_g[:, cur, None]])
    Aeq3 = np.hstack([y_b,
                      np.zeros((s2, m)),
                      np.zeros((s2, s1)),
                      np.identity(s2),
                      -y_b[:, cur, None]])
    Aeq4 = np.hstack([np.zeros(n),
                      np.zeros(m),
                      1/((s1+s2)*(y_g[:, cur])),
                      1/((s1+s2)*(y_b[:, cur])),
                      np.array([1])]).reshape(1, -1)
    if (rts == 1):
        Aeq5 = np.hstack([np.ones(n),
                          np.zeros((m+s1+s2)),
                          np.array([-1])]).reshape(1, -1)
        Aeq = np.vstack([Aeq1, Aeq2, Aeq3, Aeq4, Aeq5])
        beq = np.concatenate(
            [np.zeros(m+s1+s2), np.array([1]), np.array([0])])
    else:
        Aeq = np.vstack([Aeq1, Aeq2, Aeq3, Aeq4])
        beq = np.concatenate([np.zeros(m+s1+s2), np.array([1])])
    bounds = tuple([(0, None) for t in range(n+s1+s2+m+1)])
    res = optimize.linprog(c=f, A_eq=Aeq, b_eq=beq, bounds=bounds)
    return res

# 含期望产出的超效率SBM
def sup_un_sbmeff(x, y_g, y_b, cur, rts=0):
    m, n = x.shape  # m为投入变量数，n为决定单元数
    s1 = y_g.shape[0]
    s2 = y_b.shape[0]
    f = np.concatenate([np.zeros(n), 1/(m*x[:, cur]),
                        np.zeros(s1+s2), np.array([1])])
    Aeq1 = np.hstack([np.zeros(n),
                      np.zeros(m),
                      -1/((s1+s2)*(y_g[:, cur])),
                      -1/((s1+s2)*(y_b[:, cur])),
                      np.array([1])]).reshape(1, -1)
    if (rts == 1):
        Aeq2 = np.hstack([np.ones(n),
                          np.zeros((m+s1+s2)),
                          np.array([-1])]).reshape(1, -1)  # 规模报酬可变增加的约束
        Aeq = np.vstack([Aeq1, Aeq2])  # 合并等式约束矩阵
        Aeq[:, cur] = 0  # 将正在计算列的系数置为0，对应约束条件中的求和不包括当前dmu
        beq = np.concatenate(
            [np.array([1]), np.array([0])])  # 增加VRS约束方程右边的0
    else:
        Aeq = Aeq1
        beq = np.array([1])
    Aub1 = np.hstack([x,
                      -np.identity(m),
                      np.zeros((m, s1+s2)),
                      -x[:, cur, None]])

    Aub2 = np.hstack([-y_g,
                      np.zeros((s1, m)),
                      -np.identity(s1),
                      np.zeros((s1, s2)),
                      y_g[:, cur, None]])

    Aub3 = np.hstack([y_b,
                      np.zeros((s2, m)),
                      np.zeros((s2, s1)),
                      -np.identity(s2),
                      -y_b[:, cur, None]])
    Aub = np.vstack([Aub1, Aub2, Aub3])
    bub = np.zeros(m+s1+s2)
    Aub[:, cur] = 0 # 将正在计算列的系数置为0，对应约束条件中的求和不包括当前dmu
    bounds = tuple([(0, None) for i in range(n+s1+s2+m+1)])
    res = optimize.linprog(c=f, A_ub=Aub, b_ub=bub,
                           A_eq=Aeq, b_eq=beq, bounds=bounds)
    return res


# 填充空值的函数
def fill_missing_values(group):
    # 获取所有年份的列
    columns = group.columns[7:]

    # 对每一列应用填充逻辑
    for column in columns:
        for i in range(len(group)):
            if pd.isnull(group.iloc[i][column]):
                # 检查前2年的数据
                if i > 0 and not pd.isnull(group.iloc[i - 1][column]):
                    group.at[i, column] = group.iloc[i - 1][column]
                elif i > 1 and not pd.isnull(group.iloc[i - 2][column]):
                    group.at[i, column] = group.iloc[i - 2][column]
                # 检查第3年的数据
                elif i > 2 and not pd.isnull(group.iloc[i - 3][column]):
                    group.at[i, column] = group.iloc[i - 3][column]
    return group

def load_data(filename, sheetName='insert_data'):
    # df = pd.read_excel(filename, sheet_name=sheetName, index_col=0)
    df = pd.read_excel(filename, sheet_name=sheetName)
    # 删除包含NaN值的行
    df_cleaned = df.dropna()
    return df_cleaned.values
    # return df.values


def process_data(data, dmus, periods, nx, ny, nb, undesirable):
    theta = []
    results = []
    for t in range(periods):
        #data[0:dmus:]
        yeardata = data[dmus * t:dmus * (t + 1), :]
        # 证券代码所在列
        dmuname = yeardata[:, 0]
        # 时间
        year = yeardata[:, 1]
        # 投入
        x = yeardata[:, 4:7]
        # 产出
        y_g = yeardata[:, 2:4]
        # print(f'x={x},{len(x)} y_g={y_g}, {len(y_g)}')
        x = x.T  # 投入变量
        y_g = y_g.T  # 产出变量
        if undesirable == 1:
            y_b = yeardata[:, nx + ny + 2:nx + ny + nb + 2].T  # 非期望产出变量

        for i in range(dmus):
            try:
                if undesirable == 1:
                    res = un_sbmeff(x=x, y_g=y_g, y_b=y_b, cur=i, rts=1)
                else:
                    print(f't={t}, year={year[i]},{t*dmus},x={x},{len(x)} y_g={y_g}, {len(y_g)}, cur={i}, rts=1')
                    res = sbmeff(x=x, y_g=y_g, cur=i, rts=1)
                results.append((dmuname[i], year[i], res.fun))
                theta.append(res.fun)  # 可选，根据需求保留效率值列表
            except Exception as e:
                print(f"发生异常: {e}")

    return results, theta


def calculate_superefficiency(data, dmus, periods, nx, ny, nb, undesirable, theta):
    suptheta = []
    for t in range(periods):
        # yeardata = data[dmus * t:dmus * (t + 1), :]
        # dmuname = yeardata[:, 0]
        # year = yeardata[0, 1]
        # x = yeardata[:, 2:nx + 2].T
        # y_g = yeardata[:, nx + 2:nx + ny + 2].T

        yeardata = data[dmus * t:dmus * (t + 1), :]

        dmuname = yeardata[:, 0]

        # 时间
        year = yeardata[:, 1]
        # 投入
        x = yeardata[:, 4:7]
        # 产出
        y_g = yeardata[:, 2:4]

        x = x.T  # 投入变量
        y_g = y_g.T  # 产出变量

        if undesirable == 1:
            y_b = yeardata[:, nx + ny + 2:nx + ny + nb + 2].T

        for i in range(dmus):
            try:
                if undesirable == 1:
                    res = sup_un_sbmeff(x=x, y_g=y_g, y_b=y_b, cur=i, rts=1)
                    if res.fun > 1:
                        suptheta.append((dmuname[i], year[i], res.fun))
                    else:
                        suptheta.append(theta[dmus * t + i])
                else:
                    res = sup_sbmeff(x=x, y_g=y_g, cur=i, rts=1)
                    if res.fun != None and res.fun > 1:
                        suptheta.append((dmuname[i], year[i], res.fun))
                    else:
                        suptheta.append((dmuname[i], year[i], theta[dmus * t + i]))
            except Exception as e:
                print(f"发生异常: {e}")

    return suptheta


if __name__ == '__main__':
    dmus = 10
    periods = 1740
    nx = 2
    ny = 2
    nb = 1
    undesirable = 0   # 1：含有非期望产出 0： 无
    rts = 1
    sup = 1
    filename = 'D:\北工大\B503\合作论文\注册制和金融效率\数据\A股\核准制上证A股-数据_normalized-0_1_filled_data.xlsx'
    output_name = 'D:\北工大\B503\合作论文\注册制和金融效率\数据\A股\核准制上证A股_sbmeff_sup.xlsx'
    data = load_data(filename, sheetName='Sheet2')
    results, theta = process_data(data, dmus, periods, nx, ny, nb, undesirable)

    if sup == 0:
        dfres = pd.DataFrame(results, columns=('dmu', '时间', '效率'))
        dfres.to_excel("sbmeff.xlsx", sheet_name='eff', index=False)
        sys.exit(0)
    # print(f'results:{results}, theta:{theta}')
    suptheta = calculate_superefficiency(data, dmus, periods, nx, ny, nb, undesirable, theta)
    # print('suptheta:', suptheta)
    dfres = pd.DataFrame(suptheta, columns=('dmu', '时间', '效率'))
    dfres.to_excel(output_name, sheet_name='eff', index=False)


